#include <lua.h>
#include <lauxlib.h>

static void mark_object(lua_State *L, lua_State *dL, const void * parent, const char * desc);

#if LUA_VERSION_NUM == 501

static void
luaL_checkversion(lua_State *L) {
	if (lua_pushthread(L) == 0) {
		luaL_error(L, "Must require in main thread");
	}
	lua_setfield(L, LUA_REGISTRYINDEX, "mainthread");
}

static void
lua_rawsetp(lua_State *L, int idx, const void *p) {
	if (idx < 0) {
		idx += lua_gettop(L) + 1;
	}
	lua_pushlightuserdata(L, (void *)p);
	lua_insert(L, -2);
	lua_rawset(L, idx);
}

/// @brief 
/// @param L 
/// @param idx 
/// @param p 
static void
lua_rawgetp(lua_State *L, int idx, const void *p) {
	if (idx < 0) {
		idx += lua_gettop(L) + 1;
	}
	lua_pushlightuserdata(L, (void *)p);
	lua_rawget(L, idx);
}

static void
lua_getuservalue(lua_State *L, int idx) {
	lua_getfenv(L, idx);
}

static void
mark_function_env(lua_State *L, lua_State *dL, const void * t) {
	lua_getfenv(L,-1);
	if (lua_istable(L,-1)) {
		mark_object(L, dL, t, "[environment]");
	} else {
		lua_pop(L,1);
	}
}

// lua 5.1 has no light c function
#define is_lightcfunction(L, idx) (0)

#else

#define mark_function_env(L,dL,t)

static int
is_lightcfunction(lua_State *L, int idx) {
	if (lua_iscfunction(L, idx)) {
		if (lua_getupvalue(L, idx, 1) == NULL) {
			return 1;
		}
		lua_pop(L, 1);
	}
	return 0;
}

#endif

#include <stdbool.h>
#include <stdio.h>
#include <string.h>

#define TABLE 1
#define FUNCTION 2
#define SOURCE 3
#define THREAD 4
#define USERDATA 5
#define MARK 6

static bool
ismarked(lua_State *dL, const void *p) {
	lua_rawgetp(dL, MARK, p); //从标记的表中查找 p为key，查找dL栈中索引为MARK的表中是否存在p
	if (lua_isnil(dL,-1)) { //不存在
		lua_pop(dL,1); //将 nil 弹出
		lua_pushboolean(dL,1); //将 true 压入栈顶
		lua_rawsetp(dL, MARK, p); //将 true 作为值，p作为键，存入dL栈中索引为MARK的表中 表示这个p已经标记过了
		return false;
	}
	lua_pop(dL,1); //将查找的结果弹出
	return true;
}

static const void *
readobject(lua_State *L, lua_State *dL, const void *parent, const char *desc) {
	int t = lua_type(L, -1);
	int tidx = 0;
	switch (t) {
	case LUA_TTABLE:
		tidx = TABLE; // dL栈中为tidx的table的索引
		break;
	case LUA_TFUNCTION:
		if (is_lightcfunction(L, -1)) {
			lua_pop(L, 1);
			return NULL;
		}
		tidx = FUNCTION;
		break;
	case LUA_TTHREAD:
		tidx = THREAD;
		break;
	case LUA_TUSERDATA:
		tidx = USERDATA;
		break;
	default:
		lua_pop(L, 1);
		return NULL;
	}

	const void * p = lua_topointer(L, -1); //将栈顶的数据转化为指针
	if (ismarked(dL, p)) { //是否标记过
		lua_rawgetp(dL, tidx, p); //查找dL栈中索引为tidx的表中是否存在p，如果存在，则将对应的table压入栈顶 对应的table是 {[desc] = parent}
		if (!lua_isnil(dL,-1)) { //存在
			lua_pushstring(dL,desc);
			lua_rawsetp(dL, -2, parent); //将desc作为键，parent作为值，存入栈顶的table中，此时 栈顶的table是 {[desc1] = parent, [desc2] = parent2},将引用p的放入同一个table中
		}
		lua_pop(dL,1); //将栈顶的table或nil弹出
		lua_pop(L,1); //将栈顶的指针弹出
		return NULL;
	}

	lua_newtable(dL); //创建新表
	lua_pushstring(dL,desc);
	lua_rawsetp(dL, -2, parent); //将desc作为键，parent作为值，存入新表
	lua_rawsetp(dL, tidx, p); //将新表作为值，p作为键，存入dL最开始创建的表中   {[p] = {[desc] = [parent]}}

	return p;
}
/// @brief 将栈上索引为index的数据转为string
/// @param L 
/// @param index 索引
/// @param buffer 字符串缓冲
/// @param size 大小
/// @return 
static const char *
keystring(lua_State *L, int index, char * buffer, size_t size) {
	int t = lua_type(L,index);
	switch (t) {
	case LUA_TSTRING:
		return lua_tostring(L,index);
	case LUA_TNUMBER:
		snprintf(buffer, size, "[%lg]",lua_tonumber(L,index));
		break;
	case LUA_TBOOLEAN:
		snprintf(buffer, size, "[%s]",lua_toboolean(L,index) ? "true" : "false");
		break;
	case LUA_TNIL:
		snprintf(buffer, size, "[nil]");
		break;
	default:
		snprintf(buffer, size, "[%s:%p]",lua_typename(L,t),lua_topointer(L,index)); //将类型和指针转为字符串
		break;
	}
	return buffer;
}

static void
mark_table(lua_State *L, lua_State *dL, const void * parent, const char * desc) {
	const void * t = readobject(L, dL, parent, desc);
	if (t == NULL)
		return;

	bool weakk = false;
	bool weakv = false;
	if (lua_getmetatable(L, -1)) { //获取元表
		lua_pushliteral(L, "__mode"); 
		lua_rawget(L, -2);  //获取 t["__mode"]值，不触发元方法
		if (lua_isstring(L,-1)) { //如果是string
			const char *mode = lua_tostring(L, -1);
			if (strchr(mode, 'k')) { //是弱key
				weakk = true;
			}
			if (strchr(mode, 'v')) { //是弱value
				weakv = true;
			}
		}
		lua_pop(L,1);

		luaL_checkstack(L, LUA_MINSTACK, NULL);
		mark_table(L, dL, t, "[metatable]"); //标记元表
	}

	lua_pushnil(L);
	while (lua_next(L, -2) != 0) { //遍历table
		if (weakv) { //如果元表中的__mode为"kv"，则不标记value
			lua_pop(L,1); 
		} else { //不是弱value
			char temp[32];
			const char * desc = keystring(L, -2, temp, sizeof(temp)); //获取key的字符串
			mark_object(L, dL, t , desc);  //标记value
		}
		if (!weakk) { //不是弱key
			lua_pushvalue(L,-1);
			mark_object(L, dL, t , "[key]"); //标记key
		}
	}

	lua_pop(L,1);
}

static void
mark_userdata(lua_State *L, lua_State *dL, const void * parent, const char *desc) {
	const void * t = readobject(L, dL, parent, desc);
	if (t == NULL)
		return;
	if (lua_getmetatable(L, -1)) { //获取元表
		mark_table(L, dL, t, "[metatable]"); //标记元表
	}

	lua_getuservalue(L,-1); //获取userdata
	if (lua_isnil(L,-1)) { //是nil
		lua_pop(L,2); //将nil和table出栈
	} else {
		mark_object(L, dL, t, "[uservalue]"); //标记userdata
		lua_pop(L,1);//将table出栈
	}
}

static void
mark_function(lua_State *L, lua_State *dL, const void * parent, const char *desc) {
	const void * t = readobject(L, dL, parent, desc);
	if (t == NULL)
		return;

	mark_function_env(L,dL,t); //标记func的env
	int i;
	for (i=1;;i++) {
		const char *name = lua_getupvalue(L,-1,i); //获取所有的上值
		if (name == NULL)
			break;
		mark_object(L, dL, t, name[0] ? name : "[upvalue]"); //标记上值
	}
	if (lua_iscfunction(L,-1)) {
		lua_pop(L,1);
	} else {
		lua_Debug ar;
		lua_getinfo(L, ">S", &ar);
		luaL_Buffer b;
		luaL_buffinit(dL, &b);
		luaL_addstring(&b, "func: ");
		luaL_addstring(&b, ar.short_src);
		char tmp[16];
		sprintf(tmp,":%d",ar.linedefined);
		luaL_addstring(&b, tmp);
		luaL_pushresult(&b);
		lua_rawsetp(dL, SOURCE, t);
	}
}

static void
mark_thread(lua_State *L, lua_State *dL, const void * parent, const char *desc) {
	const void * t = readobject(L, dL, parent, desc);
	if (t == NULL)
		return;
	int level = 0;
	lua_State *cL = lua_tothread(L,-1);
	if (cL == L) {
		level = 1;
	} else {
		// mark stack
		int top = lua_gettop(cL);
		luaL_checkstack(cL, 1, NULL);
		int i;
		char tmp[16];
		for (i=0;i<top;i++) {
			lua_pushvalue(cL, i+1);
			sprintf(tmp, "[%d]", i+1);
			mark_object(cL, dL, cL, tmp);
		}
	}
	lua_Debug ar;
	luaL_Buffer b;
	luaL_buffinit(dL, &b);
	while (lua_getstack(cL, level, &ar)) {
		char tmp[128];
		lua_getinfo(cL, "Sl", &ar);
		luaL_addstring(&b, ar.short_src);
		if (ar.currentline >=0) {
			char tmp[16];
			sprintf(tmp,":%d ",ar.currentline);
			luaL_addstring(&b, tmp);
		}

		int i,j;
		// read locals (1,2 ...) and varargs (-1, -2, ...)
		for (j=1;j>=-1;j-=2) {
			for (i=j;;i+=j) {
				const char * name = lua_getlocal(cL, &ar, i);
				if (name == NULL)
					break;
				snprintf(tmp, sizeof(tmp), "%s : %s:%d",name,ar.short_src,ar.currentline);
				mark_object(cL, dL, t, tmp);
			}
		}

		++level;
	}
	luaL_addstring(&b, "thread: ");
	luaL_pushresult(&b);
	lua_rawsetp(dL, SOURCE, t);
	lua_pop(L,1);
}

static void 
mark_object(lua_State *L, lua_State *dL, const void * parent, const char *desc) {
	luaL_checkstack(L, LUA_MINSTACK, NULL);
	int t = lua_type(L, -1); //获取栈顶value的类型
	switch (t) {
	case LUA_TTABLE:
		mark_table(L, dL, parent, desc); //标记table
		break;
	case LUA_TUSERDATA:
		mark_userdata(L, dL, parent, desc); //标记userdata
		break;
	case LUA_TFUNCTION:
		mark_function(L, dL, parent, desc);//标记func
		break;
	case LUA_TTHREAD:
		mark_thread(L, dL, parent, desc);//标记协程
		break;
	default:
		lua_pop(L,1); //将栈顶数据出栈
		break;
	}
}

static int
count_table(lua_State *L, int idx) { //计算table的大小
	int n = 0;
	lua_pushnil(L);
	while (lua_next(L, idx) != 0) {
		++n;
		lua_pop(L,1);
	}
	return n;
}
/// @brief 生成table描述
/// @param dL 
/// @param b 
/// @param parent
/// @param desc 
static void
gen_table_desc(lua_State *dL, luaL_Buffer *b, const void * parent, const char *desc) {
	char tmp[32];
	size_t l = snprintf(tmp, sizeof(tmp), "%p : ",parent);
	luaL_addlstring(b, tmp, l);
	luaL_addstring(b, desc);
	luaL_addchar(b, '\n');
}
/// @brief 描述
/// @param L lua栈
/// @param dL 新创建的lua栈
/// @param idx 栈上table的索引
/// @param typename 对应table的名称
static void
pdesc(lua_State *L, lua_State *dL, int idx, const char * typename) {
	lua_pushnil(dL);
	while (lua_next(dL, idx) != 0) {
		luaL_Buffer b;
		luaL_buffinit(L, &b);//初始化字符串缓存
		const void * key = lua_touserdata(dL, -2);
		if (idx == FUNCTION) {
			lua_rawgetp(dL, SOURCE, key);
			if (lua_isnil(dL, -1)) {
				luaL_addstring(&b,"cfunction\n");
			} else {
				size_t l = 0;
				const char * s = lua_tolstring(dL, -1, &l);
				luaL_addlstring(&b,s,l);
				luaL_addchar(&b,'\n');
			}
			lua_pop(dL, 1);
		} else if (idx == THREAD) {
			lua_rawgetp(dL, SOURCE, key);
			size_t l = 0;
			const char * s = lua_tolstring(dL, -1, &l);
			luaL_addlstring(&b,s,l);
			luaL_addchar(&b,'\n');
			lua_pop(dL, 1);
		} else {
			luaL_addstring(&b, typename);
			luaL_addchar(&b,'\n');
		}
		lua_pushnil(dL);
		while (lua_next(dL, -2) != 0) {
			const void * parent = lua_touserdata(dL,-2);
			const char * desc = luaL_checkstring(dL,-1);
			printf("table_desc->desc:%s,parent:%p \n",desc, parent);
			gen_table_desc(dL, &b, parent, desc);
			lua_pop(dL,1);
		}
		luaL_pushresult(&b);
		lua_rawsetp(L, -2, key);
		lua_pop(dL,1);
	}
}

static void
gen_result(lua_State *L, lua_State *dL) {
	int count = 0;
	count += count_table(dL, TABLE); //计算table的大小
	count += count_table(dL, FUNCTION);
	count += count_table(dL, USERDATA);
	count += count_table(dL, THREAD);
	lua_createtable(L, 0, count); //创建一个table count 表示hash部分的大小
	pdesc(L, dL, TABLE, "table");
	pdesc(L, dL, USERDATA, "userdata");
	pdesc(L, dL, FUNCTION, "function");
	pdesc(L, dL, THREAD, "thread");
}

static int
snapshot(lua_State *L) {
	int i;
	lua_State *dL = luaL_newstate();
	for (i=0;i<MARK;i++) {
		lua_newtable(dL); //创建多个table 
	}
	lua_pushvalue(L, LUA_REGISTRYINDEX); //将注册表压栈
	mark_table(L, dL, NULL, "[registry]");
	gen_result(L, dL);
	lua_close(dL);
	return 1;
}

// int
// luaopen_snapshot(lua_State *L) {
// 	luaL_checkversion(L);
// 	lua_pushcfunction(L, snapshot);
// 	return 1;
// }

void
luaopen_snapshot(lua_State *L) {
	luaL_checkversion(L);
	lua_pushcfunction(L, snapshot);
	lua_setglobal(L, "snapshot");
	// return 1;
}